#!/usr/bin/env node

import express, { Request, Response } from "express";
import { randomUUID } from "node:crypto";
import { z } from "zod";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";

// 创建Express应用
const app = express();
app.use(express.json());

// 存储传输实例的映射，按会话ID索引
const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};

// 创建MCP服务器实例的函数
function createMcpServer(): McpServer {
  const server = new McpServer({
    name: "wjb-http-mcp-server",
    version: "1.0.0",
  });

  // 注册一个简单的计算器工具
  server.registerTool(
    "calculator",
    {
      title: "Calculator",
      description: "Perform basic mathematical calculations",
      inputSchema: {
        operation: z.enum(["add", "subtract", "multiply", "divide"]).describe("The operation to perform"),
        a: z.number().describe("The first number"),
        b: z.number().describe("The second number"),
      },
    },
    async ({ operation, a, b }) => {
      let result: number;
      switch (operation) {
        case "add":
          result = a + b;
          break;
        case "subtract":
          result = a - b;
          break;
        case "multiply":
          result = a * b;
          break;
        case "divide":
          if (b === 0) {
            throw new Error("Division by zero is not allowed");
          }
          result = a / b;
          break;
        default:
          throw new Error(`Unknown operation: ${operation}`);
      }
      return {
        content: [
          {
            type: "text",
            text: `${a} ${operation} ${b} = ${result}`,
          },
        ],
      };
    }
  );

  // 注册一个回声工具
  server.registerTool(
    "echo",
    {
      title: "Echo Tool",
      description: "Echo back the input text",
      inputSchema: {
        text: z.string().describe("Text to echo back"),
      },
    },
    async ({ text }) => {
      return {
        content: [
          {
            type: "text",
            text: `Echo: ${text}`,
          },
        ],
      };
    }
  );

  // 注册一个系统信息资源
  server.registerResource(
    "system-info",
    "system://info",
    {
      title: "System Information",
      description: "Basic system information",
      mimeType: "text/plain",
    },
    async (uri) => {
      const info = {
        platform: process.platform,
        arch: process.arch,
        nodeVersion: process.version,
        uptime: process.uptime(),
        timestamp: new Date().toISOString(),
      };
      return {
        contents: [
          {
            uri: uri.href,
            text: JSON.stringify(info, null, 2),
          },
        ],
      };
    }
  );

  // 注册一个时间资源
  server.registerResource(
    "current-time",
    "time://current",
    {
      title: "Current Time",
      description: "Get the current date and time",
      mimeType: "text/plain",
    },
    async (uri) => {
      const now = new Date();
      return {
        contents: [
          {
            uri: uri.href,
            text: `Current time: ${now.toISOString()}\nLocal time: ${now.toLocaleString()}`,
          },
        ],
      };
    }
  );

  return server;
}

// 处理POST请求（客户端到服务器的通信）
app.post("/mcp", async (req: Request, res: Response) => {
  try {
    const sessionId = req.headers["mcp-session-id"] as string;
    
    // 检查是否为初始化请求
    if (isInitializeRequest(req.body)) {
      // 为新会话创建新的服务器和传输实例
      const newSessionId = randomUUID();
      const server = createMcpServer();
      const transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => newSessionId,
      });
      
      // 存储传输实例
      transports[newSessionId] = transport;
      
      // 连接服务器和传输
      await server.connect(transport);
      
      // 处理请求
      await transport.handleRequest(req, res, req.body);
      
      console.log(`New session created: ${newSessionId}`);
    } else {
      // 使用现有会话
      if (!sessionId || !transports[sessionId]) {
        return res.status(400).json({
          jsonrpc: "2.0",
          error: {
            code: -32600,
            message: "Invalid session ID or session not found",
          },
        });
      }
      
      const transport = transports[sessionId];
      await transport.handleRequest(req, res, req.body);
    }
  } catch (error) {
    console.error("Error handling MCP request:", error);
    if (!res.headersSent) {
      res.status(500).json({
        jsonrpc: "2.0",
        error: {
          code: -32603,
          message: "Internal error",
        },
      });
    }
  }
});

// 处理GET请求（服务器发送事件）
app.get("/mcp", async (req: Request, res: Response) => {
  try {
    const sessionId = req.headers["mcp-session-id"] as string;
    
    if (!sessionId || !transports[sessionId]) {
      return res.status(400).json({
        error: "Invalid session ID or session not found",
      });
    }
    
    const transport = transports[sessionId];
    await transport.handleRequest(req, res);
  } catch (error) {
    console.error("Error handling SSE request:", error);
    if (!res.headersSent) {
      res.status(500).json({
        error: "Internal error",
      });
    }
  }
});

// 处理DELETE请求（终止会话）
app.delete("/mcp", async (req: Request, res: Response) => {
  try {
    const sessionId = req.headers["mcp-session-id"] as string;
    
    if (!sessionId || !transports[sessionId]) {
      return res.status(400).json({
        error: "Invalid session ID or session not found",
      });
    }
    
    const transport = transports[sessionId];
    transport.close();
    delete transports[sessionId];
    
    console.log(`Session terminated: ${sessionId}`);
    res.status(200).json({ message: "Session terminated" });
  } catch (error) {
    console.error("Error terminating session:", error);
    if (!res.headersSent) {
      res.status(500).json({
        error: "Internal error",
      });
    }
  }
});

// 健康检查端点
app.get("/health", (req: Request, res: Response) => {
  res.json({
    status: "healthy",
    timestamp: new Date().toISOString(),
    activeSessions: Object.keys(transports).length,
  });
});

// 根路径信息
app.get("/", (req: Request, res: Response) => {
  res.json({
    name: "WJB HTTP MCP Server",
    version: "1.0.0",
    description: "HTTP MCP Server using TypeScript SDK",
    endpoints: {
      mcp: "/mcp",
      health: "/health",
    },
    activeSessions: Object.keys(transports).length,
  });
});

// 启动服务器
const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {
  console.log(`🚀 HTTP MCP Server is running on port ${PORT}`);
  console.log(`📡 MCP endpoint: http://localhost:${PORT}/mcp`);
  console.log(`🏥 Health check: http://localhost:${PORT}/health`);
  console.log(`📖 Server info: http://localhost:${PORT}/`);
});

// 优雅关闭
process.on("SIGINT", () => {
  console.log("\n🛑 Shutting down server...");
  // 关闭所有活动的传输
  Object.values(transports).forEach(transport => transport.close());
  process.exit(0);
});

process.on("SIGTERM", () => {
  console.log("\n🛑 Shutting down server...");
  // 关闭所有活动的传输
  Object.values(transports).forEach(transport => transport.close());
  process.exit(0);
});